﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Windows.Forms;

// 客户端会话信息类
public class ClientSession
{
    public string SessionId { get; }
    public Socket ClientSocket { get; }
    public IPEndPoint RemoteEndPoint { get; }
    public DateTime ConnectedTime { get; }

    public ClientSession(Socket socket)
    {
        SessionId = Guid.NewGuid().ToString("N");
        ClientSocket = socket;
        RemoteEndPoint = socket.RemoteEndPoint as IPEndPoint;
        ConnectedTime = DateTime.Now;
    }
}

public class DataReceivedEventArgs : EventArgs
{
    public ClientSession Session { get; }
    public string Data { get; }
    public string Response { get; set; }
    public bool ShouldRespond => !string.IsNullOrEmpty(Response);

    public DataReceivedEventArgs(ClientSession session, string data)
    {
        Session = session;
        Data = data;
    }
}

// 日志事件参数（仍保留，供UI订阅）
public class LogEventArgs : EventArgs
{
    public string Message { get; }
    public DateTime Time { get; }

    public LogEventArgs(string message)
    {
        Message = message;
        Time = DateTime.Now;
    }
}

public class EnhancedSocketServer
{
    private readonly Socket _listenSocket;
    private readonly SocketAsyncEventArgsPool _eventArgsPool;
    private readonly ConcurrentDictionary<string, ClientSession> _sessions = new ConcurrentDictionary<string, ClientSession>();
    private readonly ConcurrentDictionary<Socket, string> _socketToSessionId = new ConcurrentDictionary<Socket, string>();

    private const int BufferSize = 1024;
    private const int MaxConnections = 1000;

    public event EventHandler<DataReceivedEventArgs> OnReceiveData;
    public event EventHandler<LogEventArgs> OnLog;  // 日志事件，供WinForm订阅

    public EnhancedSocketServer(int port)
    {
        _listenSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        _listenSocket.Bind(new IPEndPoint(IPAddress.Any, port));
        _listenSocket.Listen(MaxConnections);

        _eventArgsPool = new SocketAsyncEventArgsPool(MaxConnections);
        for (int i = 0; i < MaxConnections; i++)
        {
            var args = new SocketAsyncEventArgs();
            args.Completed += OnIOCompleted;
            args.SetBuffer(new byte[BufferSize], 0, BufferSize);
            _eventArgsPool.Push(args);
        }

        // 使用Debug输出日志
        string logMsg = $"服务器已启动，监听端口: {port}";
        Debug.WriteLine(logMsg);
        OnLog?.Invoke(this, new LogEventArgs(logMsg));
    }

    public void Start()
    {
        StartAccept(null);
    }

    // 获取所有客户端会话
    public List<ClientSession> GetAllSessions() => _sessions.Values.ToList();

    // 根据SessionId获取会话
    public ClientSession GetSessionById(string sessionId)
    {
        _sessions.TryGetValue(sessionId, out var session);
        return session;
    }

    // 向指定会话发送数据
    public bool SendToSession(string sessionId, string data)
    {
        var session = GetSessionById(sessionId);
        return session != null && SendToSocket(session.ClientSocket, data);
    }

    // 向指定Socket发送数据
    public bool SendToSocket(Socket socket, string data)
    {
        if (socket == null || !socket.Connected || !_socketToSessionId.ContainsKey(socket))
            return false;

        try
        {
            var args = _eventArgsPool.Pop();
            if (args == null)
                return false;

            byte[] sendData = Encoding.UTF8.GetBytes(data);
            args.UserToken = socket;
            args.SetBuffer(sendData, 0, sendData.Length);

            bool willRaiseEvent = socket.SendAsync(args);
            if (!willRaiseEvent)
            {
                ProcessSend(args);
            }
            return true;
        }
        catch (Exception ex)
        {
            string logMsg = $"发送数据错误: {ex.Message}";
            Debug.WriteLine(logMsg);
            OnLog?.Invoke(this, new LogEventArgs(logMsg));
            return false;
        }
    }

    // 广播数据给所有客户端
    public void Broadcast(string data)
    {
        foreach (var session in _sessions.Values.ToArray())
        {
            SendToSocket(session.ClientSocket, data);
        }
    }

    private void StartAccept(SocketAsyncEventArgs acceptEventArgs)
    {
        if (acceptEventArgs == null)
        {
            acceptEventArgs = new SocketAsyncEventArgs();
            acceptEventArgs.Completed += OnAcceptCompleted;
        }
        else
        {
            acceptEventArgs.AcceptSocket = null;
        }

        bool willRaiseEvent = _listenSocket.AcceptAsync(acceptEventArgs);
        if (!willRaiseEvent)
        {
            ProcessAccept(acceptEventArgs);
        }
    }

    private void OnAcceptCompleted(object sender, SocketAsyncEventArgs e)
    {
        ProcessAccept(e);
    }

    private void ProcessAccept(SocketAsyncEventArgs e)
    {
        if (e.SocketError == SocketError.Success)
        {
            var clientSocket = e.AcceptSocket;
            var session = new ClientSession(clientSocket);

            _sessions.TryAdd(session.SessionId, session);
            _socketToSessionId.TryAdd(clientSocket, session.SessionId);

            string logMsg = $"客户端连接 - SessionId: {session.SessionId}, 地址: {session.RemoteEndPoint?.Address}";
            Debug.WriteLine(logMsg);
            OnLog?.Invoke(this, new LogEventArgs(logMsg));

            SocketAsyncEventArgs readEventArgs = _eventArgsPool.Pop();
            if (readEventArgs != null)
            {
                readEventArgs.UserToken = clientSocket;

                bool willRaiseEvent = clientSocket.ReceiveAsync(readEventArgs);
                if (!willRaiseEvent)
                {
                    ProcessReceive(readEventArgs);
                }
            }
            else
            {
                CloseSession(session.SessionId);
                logMsg = "连接池已满，拒绝新连接";
                Debug.WriteLine(logMsg);
                OnLog?.Invoke(this, new LogEventArgs(logMsg));
            }
        }

        StartAccept(e);
    }

    private void OnIOCompleted(object sender, SocketAsyncEventArgs e)
    {
        switch (e.LastOperation)
        {
            case SocketAsyncOperation.Receive:
                ProcessReceive(e);
                break;
            case SocketAsyncOperation.Send:
                ProcessSend(e);
                break;
            default:
                throw new ArgumentException("未知的操作: " + e.LastOperation);
        }
    }

    private void ProcessReceive(SocketAsyncEventArgs e)
    {
        Socket clientSocket = (Socket)e.UserToken;

        if (e.BytesTransferred == 0 || e.SocketError != SocketError.Success)
        {
            CloseSocket(clientSocket);
            return;
        }

        if (_socketToSessionId.TryGetValue(clientSocket, out string sessionId) &&
            _sessions.TryGetValue(sessionId, out ClientSession session))
        {
            string receivedData = Encoding.UTF8.GetString(e.Buffer, e.Offset, e.BytesTransferred);
            string logMsg = $"收到数据 - SessionId: {sessionId}, 数据: {receivedData}";
            Debug.WriteLine(logMsg);
            OnLog?.Invoke(this, new LogEventArgs(logMsg));

            var eventArgs = new DataReceivedEventArgs(session, receivedData);
            OnReceiveData?.Invoke(this, eventArgs);

            if (eventArgs.ShouldRespond)
            {
                byte[] sendData = Encoding.UTF8.GetBytes(eventArgs.Response);
                e.SetBuffer(sendData, 0, sendData.Length);

                bool willRaiseEvent = clientSocket.SendAsync(e);
                if (!willRaiseEvent)
                {
                    ProcessSend(e);
                }
            }
            else
            {
                e.SetBuffer(new byte[BufferSize], 0, BufferSize);
                bool willRaiseEvent = clientSocket.ReceiveAsync(e);
                if (!willRaiseEvent)
                {
                    ProcessReceive(e);
                }
            }
        }
        else
        {
            CloseSocket(clientSocket);
        }
    }

    private void ProcessSend(SocketAsyncEventArgs e)
    {
        if (e.SocketError != SocketError.Success)
        {
            Socket clientSocket = (Socket)e.UserToken;
            CloseSocket(clientSocket);
            return;
        }

        e.SetBuffer(new byte[BufferSize], 0, BufferSize);
        Socket socket = (Socket)e.UserToken;

        bool willRaiseEvent = socket.ReceiveAsync(e);
        if (!willRaiseEvent)
        {
            ProcessReceive(e);
        }
    }

    private void CloseSocket(Socket socket)
    {
        if (socket == null) return;

        if (_socketToSessionId.TryRemove(socket, out string sessionId))
        {
            CloseSession(sessionId);
        }
        else
        {
            try
            {
                socket.Shutdown(SocketShutdown.Both);
            }
            catch { }
            finally
            {
                socket.Dispose();
            }
        }
    }

    public void CloseSession(string sessionId)
    {
        if (_sessions.TryRemove(sessionId, out ClientSession session))
        {
            _socketToSessionId.TryRemove(session.ClientSocket, out _);

            string clientAddress = session.RemoteEndPoint?.Address.ToString() ?? "未知地址";
            string logMsg = $"客户端断开 - SessionId: {sessionId}, 地址: {clientAddress}";
            Debug.WriteLine(logMsg);
            OnLog?.Invoke(this, new LogEventArgs(logMsg));

            try
            {
                session.ClientSocket.Shutdown(SocketShutdown.Both);
            }
            catch { }
            finally
            {
                session.ClientSocket.Dispose();
            }
        }
    }
}

public class SocketAsyncEventArgsPool
{
    private readonly Stack<SocketAsyncEventArgs> _pool;
    private readonly int _capacity;

    public SocketAsyncEventArgsPool(int capacity)
    {
        _capacity = capacity;
        _pool = new Stack<SocketAsyncEventArgs>(capacity);
    }

    public void Push(SocketAsyncEventArgs item)
    {
        if (item == null)
            throw new ArgumentNullException(nameof(item));

        lock (_pool)
        {
            _pool.Push(item);
        }
    }

    public SocketAsyncEventArgs Pop()
    {
        lock (_pool)
        {
            return _pool.Count > 0 ? _pool.Pop() : null;
        }
    }

    public int Count => _pool.Count;
}
